import os
import matplotlib.pyplot as plt
import seaborn as sns
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split

plt.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'SimSun', 'FangSong', 'KaiTi']  # 指定一系列备选的中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号'-'显示为方块的问题


def ana_data(data):
    # print(data.info())
    # 1.总体分布
    a_data = data.copy()
    fig = plt.figure(figsize=(10, 10))
    ax1 = fig.add_subplot()
    status_count = a_data['Attrition'].value_counts()
    print(status_count)
    ax1.pie(status_count, labels=["未离职", "离职"], autopct='%.2f%%')
    ax1.set_title("离职总体分布情况")

    numeric_features = ['Age', 'Department', 'DistanceFromHome', 'Education', 'EnvironmentSatisfaction', 'Gender',
                        'JobInvolvement', 'JobLevel', 'JobRole', 'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome',
                        'OverTime',
                        'PercentSalaryHike', 'StockOptionLevel', 'WorkLifeBalance', 'YearsAtCompany',
                        'YearsSinceLastPromotion',
                        'RelationshipSatisfaction']
    # numeric_features.remove('Attrition')
    plt.figure(figsize=(16, 12))
    for i, col in enumerate(numeric_features[:]):
        plt.subplot(4, 5, i + 1)
        sns.histplot(a_data[col], kde=True, color='green')
        plt.title(f'{col}分布')
    plt.tight_layout()
    plt.savefig('特征的分布.png')
    plt.show()

    # 2.相关性分布
    # 2.1 数值型特征与离职率相关性
    numeric_features_num = ['Age', 'DistanceFromHome', 'Education', 'EnvironmentSatisfaction', 'JobInvolvement',
                            'JobLevel', 'JobSatisfaction', 'MonthlyIncome', 'PercentSalaryHike', 'StockOptionLevel',
                            'WorkLifeBalance', 'YearsAtCompany', 'YearsSinceLastPromotion',
                            'RelationshipSatisfaction']
    plt.figure(figsize=(16, 12))
    for i, col in enumerate(numeric_features_num[:]):
        plt.subplot(4, 5, i + 1)
        sns.kdeplot(data=a_data, x=col, hue="Attrition", fill=True, common_norm=False, alpha=0.6,
                    palette=['skyblue', 'orange'])
        plt.title(f'{col}与离职率的相关性分布')
        plt.legend(title='Attrition', labels=['离职', '未离职'])
    plt.tight_layout()
    plt.savefig('数值型特征与离职率相关性.png')
    plt.show()

    # 2.2 分类型特征与离职率相关性
    categorical_features = ['BusinessTravel', 'Department', 'EducationField', 'Gender', 'JobRole', 'MaritalStatus',
                            'Over18', 'OverTime']
    plt.figure(figsize=(16, 15))  # 调整图像大小以适应更多子图
    for i, col in enumerate(categorical_features):
        plt.subplot(3, 3, i + 1)  # 修改为3行3列，最多可容纳9个子图
        sns.countplot(data=a_data, x=col, hue='Attrition')
        plt.title(f'{col}与离职率的相关性分布')
        plt.legend(title='Attrition', labels=['离职', '未离职'])
    plt.tight_layout()
    plt.savefig('分类型特征与离职率相关性.png')
    plt.show()


if __name__ == '__main__':
    train_df = pd.read_csv('../data/raw/train.csv')
    ana_data(train_df)
